from matplotlib.figure import Figure
from dataclasses import dataclass, asdict, field
from matplotlib import pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.axes import Axes
from typing import Literal
import numpy as np
import re
import os
from typing import Callable
from cartopy import crs as ccrs
from cartopy.crs import Projection
from rich.progress import Progress
from xarray import DataArray

class ExitLoop(Exception):
    pass

@dataclass
class SuperRcParams:
    """
    rcParams 的配置类型，使用驼峰命名法填 rcParams 的配置，
    然后使用 apply() 方法应用到 plt.rcParams 中。

    例如设置 rcParams['font.size'] = 10.5，可令 fontSize = 10.5
    """
    def _camel_to_dot(self, name):
        """驼峰命名转换为点分割命名"""
        s1 = re.sub('(.)([A-Z][a-z]+)', r'\1.\2', name)
        return re.sub('([a-z0-9])([A-Z])', r'\1.\2', s1).lower()

    def apply(self):
        """将对象中的属性应用到plt.rcParams中
        """
        for key, value in asdict(self).items():
            key = self._camel_to_dot(key)
            plt.rcParams[key] = value


@dataclass
class StandardRcParams(SuperRcParams):
    """用于A4印刷的标准配置"""
    fontSize: float = 10.5

    # 边框线宽
    axesLinewidth: float = 0.8

    # 刻度线宽
    xtickMajorWidth: float = 0.8
    ytickMajorWidth: float = 0.8
    xtickMinorWidth: float = 0.6
    ytickMinorWidth: float = 0.6

    # 刻度线长度
    xtickMajorSize: float = 4
    ytickMajorSize: float = 4
    xtickMinorSize: float = 2.5
    ytickMinorSize: float = 2.5

    # 刻度线标签间隔
    xtickMajorPad: float = 3
    ytickMajorPad: float = 3

    # 刻度线标签大小
    xtickLabelsize: float = fontSize
    ytickLabelsize: float = fontSize


@dataclass
class MapConfig:
    coastline: bool = True
    coastlineWidth: float = 0.8
    coastlineColor: str = 'k'
    coastlineKwargs: dict = field(default_factory=dict)

    fillLand: bool = False
    landColor: str = 'lightgray'

    gridline: bool = True
    gridlineStyle: str = '--'
    gridlineWidth: float = 0.35
    gridlineColor: str = 'gray'
    gridlineKwargs: dict = field(default_factory=dict)

    lonMajorInterval: float = 20
    lonMinorInterval: float = 10
    latMajorInterval: float = 20
    latMinorInterval: float = 10


from cartopy.mpl.geoaxes import GeoAxes
from numpy import ndarray

class MapTools():
    def __init__(self, ax: GeoAxes, mapConfig: MapConfig):
        self.ax = ax
        self.mapConfig = mapConfig

    def mapTicks(self, lonTicks: ndarray | None=None, 
                 latTicks: ndarray | None=None,
                 lon: ndarray | None=None, lat: ndarray | None=None,
                 drawLonLabel: bool=True, drawLatLabel: bool=True,
                 drawMinorTick: bool=True, drawGridline: bool=True,
                 useAxesTicks: bool=True, mapConfig: MapConfig=None,
                 drawMinorGridline: bool=False,
                 ignoreFristLonTickLabel: bool=False):

        # 如果没有指定 mapConfig，则使用默认配置
        if mapConfig is None:
            c = self.mapConfig
        else:
            c = mapConfig

        # 绘制海岸线
        if c.coastline:
            self.ax.coastlines(linewidth=c.coastlineWidth, color=c.coastlineColor,
                               **c.coastlineKwargs)
        
        # 填充陆地
        if c.fillLand:
            import cartopy.feature as cfeature
            self.ax.add_feature(cfeature.LAND, facecolor=c.landColor)

        from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
        from matplotlib.ticker import MultipleLocator

        # 如果不指定刻度，则用主刻度间隔生成
        extent = self.ax.get_extent()
        if lonTicks is None:
            allLonTicks = np.arange(-180, 361, c.lonMajorInterval)
            if lon is not None:
                lonTicks = allLonTicks[(lonTicks >= np.min(lon)) & 
                                        (lonTicks <= np.max(lon))]
            else:
                lonTicks = allLonTicks[(allLonTicks >= extent[0]) &
                                        (allLonTicks <= extent[1])]
        
        if latTicks is None:
            allLatTicks = np.arange(-80, 81, c.latMajorInterval)
            if lat is not None:
                latTicks = allLatTicks[(latTicks >= np.min(lat)) & 
                                        (latTicks <= np.max(lat))]
            else:
                latTicks = allLatTicks[(allLatTicks >= extent[2]) &
                                        (allLatTicks <= extent[3])]

        if ignoreFristLonTickLabel:
            lonTicks = lonTicks[1:]
        
        self.ax.xaxis.set_major_formatter(LongitudeFormatter())
        self.ax.yaxis.set_major_formatter(LatitudeFormatter())
        self.ax.set_xticks(lonTicks, crs=self.ax.projection)
        self.ax.set_yticks(latTicks, crs=self.ax.projection)

        if not drawLonLabel:
            self.ax.set_xticklabels([])
        if not drawLatLabel:
            self.ax.set_yticklabels([])

        if drawMinorTick:
            self.ax.xaxis.set_minor_locator(MultipleLocator(c.lonMinorInterval))
            self.ax.yaxis.set_minor_locator(MultipleLocator(c.latMinorInterval))

        if drawGridline:
            if drawMinorGridline:
                lonTicks = np.arange(lonTicks[0] - c.lonMinorInterval, 
                                     lonTicks[-1] + 2 * c.lonMajorInterval, 
                                     c.lonMinorInterval)
                latTicks = np.arange(latTicks[0] - c.latMinorInterval,
                                     latTicks[-1] + 2 * c.latMajorInterval,
                                     c.latMinorInterval)

            gl = self.ax.gridlines(crs=self.ax.projection, 
                                   draw_labels=not useAxesTicks,
                                   xlocs=lonTicks, ylocs=latTicks, 
                                   linewidth=c.gridlineWidth, 
                                   color=c.gridlineColor,
                                   linestyle=c.gridlineStyle,
                                   **c.gridlineKwargs, )
            gl.top_labels = False
            gl.right_labels = False

            if useAxesTicks:
                gl.left_labels = False
                gl.bottom_labels = False
            else:
                gl.left_labels = drawLatLabel
                gl.bottom_labels = drawLonLabel


class MplTools():
    def __init__(self, ax: Axes):
        self.ax = ax

    def quiver(self, x: DataArray | np.ndarray, y: DataArray | np.ndarray,
                u: DataArray | np.ndarray, v: DataArray | np.ndarray, 
                proj: ccrs.Projection = ccrs.PlateCarree(),
                scale: float = 1, units: str = 'xy', width: float = 0.25,
                lowRes: int = 4, zorder:int = 5, **kwargs):
        
        if isinstance(x, DataArray):
            x = x.values
        
        if isinstance(y, DataArray):
            y = y.values
        
        if isinstance(u, DataArray):
            u = u.values

        if isinstance(v, DataArray):
            v = v.values
        
        q = self.ax.quiver(x[::lowRes], y[::lowRes], 
                    u[::lowRes, ::lowRes], v[::lowRes, ::lowRes], 
                    transform=proj, scale=scale, units=units,
                    width=width, zorder=zorder, **kwargs)
        return q

    def draw_box_by_center(self, centerX: float, centerY: float,
                           width: float, height: float, seg: int=4, **kwargs):
        """根据中心点绘制矩形框"""
        x0 = centerX - width/2
        x1 = centerX + width/2
        y0 = centerY - height/2
        y1 = centerY + height/2
        x = np.concatenate([np.linspace(x0, x1, seg), np.linspace(x1, x1, seg),
                            np.linspace(x1, x0, seg), np.linspace(x0, x0, seg)])
        y = np.concatenate([np.linspace(y0, y0, seg), np.linspace(y0, y1, seg),
                            np.linspace(y1, y1, seg), np.linspace(y1, y0, seg)])
        c = self.ax.plot(x, y, **kwargs)
        return c


@dataclass
class SubplotParams:
    """子图绘图的回调参数"""
    row: int
    col: int
    index: int
    """以行为主的索引"""

    colIndex: int
    """以列为主的索引"""

    ax: Axes | GeoAxes
    fig: "SuperFigure"
    rows: int = 0
    cols: int = 0
    mapTools: MapTools | None = None
    mplTools: MplTools | None = None



    @property
    def isLeft(self) -> bool:
        return self.col == 0
    
    @property
    def isRight(self) -> bool:
        return self.col == self.cols - 1
    
    @property
    def isTop(self) -> bool:
        return self.row == 0
    
    @property
    def isBottom(self) -> bool:
        return self.row == self.rows - 1

    def label(self, text: str, loc: Literal['lt', 'lb', 'rt', 'rb']='lt', 
              backGroundColor: bool | str= 'white', bgPad: float=0.5,
              zorder=100, offset: tuple[float, float]=(0.01, 0.01), **kwargs):
              
        """在子图上添加标签

        Args:
            text (str): 文本
            loc (Literal['lt', 'lb', 'rt', 'rb'], optional): 位置. Defaults to 'lt'.
            offset (tuple[float, float], optional): 偏移. Defaults to (0, 0).
        """
        if loc == 'lt':
            x = offset[0]
            y = 1 - offset[1]
            ha = 'left'
            va = 'top'
        elif loc == 'lb':
            x = offset[0]
            y = offset[1]
            ha = 'left'
            va = 'bottom'
        elif loc == 'rt':
            x = 1- offset[0]
            y = 1 - offset[1]
            ha = 'right'
            va = 'top'
        elif loc == 'rb':
            x = 1 - offset[0]
            y = offset[1]
            ha = 'right'
            va = 'bottom'

        if backGroundColor:
            self.ax.text(x, y, text, 
                         ha=ha, va=va, transform=self.ax.transAxes, 
                         bbox=dict(facecolor='white', edgecolor='none', pad=bgPad), zorder=zorder, **kwargs)
        else:
            self.ax.text(x, y, text, 
                     ha=ha, va=va, transform=self.ax.transAxes, 
                     zorder=zorder, **kwargs)
        
    def log_y(self, yticks: list[float] | ndarray, 
              yticklabels: list[str] | ndarray =None,
              minorInterval: float=100):
        if yticklabels is None:
            yticklabels = yticks
            
        self.ax.set_yscale('log')
        self.ax.set_yticks(yticks)
        self.ax.set_yticklabels(yticklabels)
        # self.ax.yaxis.set_minor_locator(plt.NullLocator())
        self.ax.yaxis.set_minor_locator(plt.MultipleLocator(minorInterval))


class SuperFigure(Figure):
    def __init__(self, a4Size: bool=False, a4PaddingInch: float=1,
                 a4HeightInch: float=5, stdConfig: bool=True,
                 colorBar: Literal[False, 'right', 'bottom'] = False,
                 colorBarAreaWidthInch: float=0.8, colorbarNumber: int=1,
                 colorBarBarWidthInch: float=0.2, colorBarMarginInch: float=0.2,
                 colorBarPaddingInch: float=0,
                  *args, **kwargs):
        """
        colorBarMarginInch 是与子图的间距，colorBarPaddingInch 是 colorbar 箭头两边的内间距 
        """
        if a4Size:
            kwargs['figsize'] = (8.27-a4PaddingInch*2, a4HeightInch)

        self.useA4Size: bool = a4Size

        self.gridSpec: GridSpec | None = None
        self.nrows: int = 1
        self.ncols: int = 1
        self.subplotBaseWidthInch: float = 2
        self.leftPt: float = 50
        self.rightPt: float = 20
        self.topPt: float = 20
        self.bottomPt: float = 50
        self.hspacePt: float = 20
        self.wspacePt: float = 20

        self.useColorbar = colorBar
        self.colorbarAreaWidthInch = colorBarAreaWidthInch
        self.colorbarNumber = colorbarNumber
        self.colorbarBarWidthInch = colorBarBarWidthInch
        self.colorbarMarginInch = colorBarMarginInch
        self.colorbarPaddingInch = colorBarPaddingInch

        self.colorbarGs = None
        
        super().__init__(*args, **kwargs)

        if stdConfig:
            StandardRcParams().apply()
            self.useFont()

        self.fontsize = plt.rcParams['font.size']


    def useFont(self, fontName: str="Helvetica", asDefault: bool=True,
                asMathText: bool=False) -> None:
        """使用自定义字体，字体文件需要放在 draw/fonts 文件夹下

        Args:
            fontName (str, optional): 字体名. Defaults to "Helvetica".
            asDefault (bool, optional): 是否作为第一字体. Defaults to True.
            asMathText (bool, optional): 使用用在公式中. Defaults to False.
        """
        from matplotlib import font_manager

        # TODO: 可添加自定义字体文件

        fontSets = {
            "Helvetica": "Helvetica.ttc",
            "Cascadia Mono": "CascadiaMono.ttf",
            "Oppo Sans": "OPPOSans-Regular.ttf",
        }
        dir = os.path.dirname(os.path.realpath(__file__)) + '/fonts/'
        folderPath = dir + fontSets[fontName]
        font_manager.fontManager.addfont(folderPath)
        if asDefault:
            plt.rcParams['font.family'] = fontName
        else:
            folderPath = '/Users/jjli/code/fonts/Helvetica.ttc'
            font_manager.fontManager.addfont(folderPath)
            plt.rcParams['font.family'] = ['Helvetica', fontName]

        if asMathText:
            plt.rcParams['mathtext.default'] = 'regular'    


    def subplots_adjust_inch(self, left: float = None, bottom: float = None, 
                             right: float = None, top: float = None) -> None:
        """调整整个图的边距，与subplots_adjust()相同， 但使用的是英寸为单位，
        即与 figsize 的单位相同。
        需要注意的是，值为边距，即内容到各个边的距离，而不是边框到原点的距离。

        Args:
            left (float, optional): 左边边距. Defaults to None.
            bottom (float, optional): 下边边距. Defaults to None.
            right (float, optional): 右边边距. Defaults to None.
            top (float, optional): 上边边距. Defaults to None.
        """
        width = self.get_figwidth()
        height = self.get_figheight()
        super().subplots_adjust(
            left=left/width if left is not None else None,
            bottom=bottom/height if bottom is not None else None,
            right=1-right/width if right is not None else None,
            top=1-top/height if top is not None else None
        )


    def subplots_adjust_pt(self, left: float = None, bottom: float = None, 
                           right: float = None, top: float = None) -> None:
        """调整整个图的边距，与subplots_adjust()相同， 但使用的是磅为单位，
        即与字体大小的单位相同。

        Args:
            left (float, optional): 左边边距. Defaults to None.
            bottom (float, optional): 下边边距. Defaults to None.
            right (float, optional): 右边边距. Defaults to None.
            top (float, optional): 上边边距. Defaults to None.
        """
        width = self.get_figwidth() * 72
        height = self.get_figheight() * 72
        super().subplots_adjust(
            left=left/width if left is not None else None,
            bottom=bottom/height if bottom is not None else None,
            right=1-right/width if right is not None else None,
            top=1-top/height if top is not None else None
        )


    def _pt2inch(self, pt: float) -> float:
        return pt/72


    def _inch2pt(self, inch: float) -> float:
        return inch*72


    def ratio_gridspec(self, nrows: int=1, ncols: int=1, 
                        subplotRatio: float=1,
                        subplotWidthInch: float=2, 
                        hspacePt: float=20,
                        wspacePt: float=20,
                        leftPt: float=50, rightPt: float=20,
                        topPt: float=20, bottomPt: float=50
                        ) -> GridSpec:
        """创建一个固定比例的网格布局，如果没有设置 a4Size，则会根据每个子图的
        宽度 subplotWidthInch 为基准，重新调整图片的大小。
        否则会根据 A4 纸的大小调整每个子图的大小。

        注意边距与间隔宽度的单位为磅，即与字体大小的单位相同，与原来的比例不同。

        Args:
            nrows (int, optional): 行数. Defaults to 1.
            ncols (int, optional): 列数. Defaults to 1.
            subplotRatio (float, optional): 宽高比，宽/高. Defaults to 1.
            subplotWidthInch (float, optional): 子图的宽度，以英寸为单位. Defaults to 2.
            hspacePt (float, optional): 子图间的纵向间隔高度，以点为单位. Defaults to 20.
            wspacePt (float, optional): 子图间的横向间隔宽度，以点为单位. Defaults to 20.
            leftPt (float, optional): 全图左边距，以点为单位. Defaults to 50.
            rightPt (float, optional): 全图右边距，以点为单位. Defaults to 20.
            topPt (float, optional): 全图上边距，以点为单位. Defaults to 20.
            bottomPt (float, optional): 全图下边距，以点为单位. Defaults to 50.
        """

        self.nrows = nrows
        self.ncols = ncols
        self.subplotBaseWidthInch = subplotWidthInch
        self.hspacePt = hspacePt
        self.wspacePt = wspacePt
        self.leftPt = leftPt
        self.rightPt = rightPt
        self.topPt = topPt
        self.bottomPt = bottomPt

        colorBarWidth = 0
        colorBarHeight = 0
        colorBarWspace = 0
        colorBarHspace = 0
        colorBarMarginR = 0
        colorBarMarginB = 0
        colorBarPadiingLR = 0
        colorBarPadiingTB = 0

        if self.useColorbar == 'right':
            colorBarWidth = self.colorbarAreaWidthInch
            colorBarHeight = 0
            colorBarWspace = self._pt2inch(wspacePt)
            colorBarMarginR = self.colorbarMarginInch
            colorBarMarginB = 0
            colorBarPadiingLR = 0
            colorBarPadiingTB = self.colorbarPaddingInch

        elif self.useColorbar == 'bottom':
            colorBarWidth = 0
            colorBarHeight = self.colorbarAreaWidthInch
            colorBarHspace = self._pt2inch(hspacePt)
            colorBarMarginR = 0
            colorBarMarginB = self.colorbarMarginInch
            colorBarPadiingLR = self.colorbarPaddingInch
            colorBarPadiingTB = 0
        
        if self.useA4Size:
            figWidth = 8.27
            # 令 subplotWidthInch 变为被动
            subplotWidthInch = (figWidth - self._pt2inch(leftPt + rightPt) - \
                self._pt2inch(wspacePt) * (ncols-1) - 
                colorBarWidth - colorBarWspace) / ncols
        else:
            figWidth = subplotWidthInch * ncols + self._pt2inch(leftPt + rightPt) \
                + self._pt2inch(wspacePt) * (ncols-1) + colorBarWidth + colorBarWspace

        figHeight = subplotWidthInch / subplotRatio * nrows + \
            colorBarHeight + colorBarHspace + \
            self._pt2inch(topPt + bottomPt) + self._pt2inch(hspacePt) * (nrows-1)
        
        self.set_size_inches(figWidth, figHeight)

        self.gridSpec = GridSpec(nrows, ncols, 
            hspace=self._pt2inch(hspacePt) / (subplotWidthInch / subplotRatio),
            wspace=self._pt2inch(wspacePt) / subplotWidthInch,
            left=self._pt2inch(leftPt)/figWidth, 
            right=1-(self._pt2inch(rightPt) + colorBarWidth + colorBarWspace + colorBarMarginR)/figWidth,
            top=1-self._pt2inch(topPt)/figHeight, 
            bottom=(self._pt2inch(bottomPt) + colorBarHeight + colorBarHspace + colorBarMarginB)/figHeight
        )

        if self.useColorbar == 'right':
            self.colorbarGs = GridSpec(self.colorbarNumber, 1,
                hspace=self._pt2inch(hspacePt) / (subplotWidthInch / subplotRatio),
                wspace=self._pt2inch(wspacePt) / subplotWidthInch,
                left=1-(self._pt2inch(rightPt) + colorBarWidth)/figWidth,
                right=1-(self._pt2inch(rightPt)+(self.colorbarAreaWidthInch-self.colorbarBarWidthInch))/figWidth,
                top=1-(self._pt2inch(topPt)+colorBarPadiingTB)/figHeight,
                bottom=(self._pt2inch(bottomPt)+colorBarPadiingTB)/figHeight
            )
        elif self.useColorbar == 'bottom':
            self.colorbarGs = GridSpec(1, self.colorbarNumber,
                hspace=self._pt2inch(hspacePt) / (subplotWidthInch / subplotRatio),
                wspace=self._pt2inch(wspacePt) / subplotWidthInch,
                left=(self._pt2inch(leftPt)+colorBarPadiingLR)/figWidth,
                right=1-(self._pt2inch(rightPt)+colorBarPadiingLR)/figWidth,
                top=(self._pt2inch(bottomPt) + colorBarHeight)/figHeight,
                bottom=(self._pt2inch(bottomPt)+(self.colorbarAreaWidthInch-self.colorbarBarWidthInch))/figHeight
            )

        return self.gridSpec

    
    def set_colorbar(self, index: int, mappable, **kwargs):
        if self.useColorbar == 'right':
            return self.colorbar(mappable, 
                                orientation='vertical',
                                 cax=self.add_subplot(self.colorbarGs[index, :]), 
                                 **kwargs)
        elif self.useColorbar == 'bottom':
            return self.colorbar(mappable, 
                                 cax=self.add_subplot(self.colorbarGs[:, index]), 
                                 orientation='horizontal', **kwargs)


    def apply_all_subplots(self, callback: Callable[[SubplotParams], None],
                           showProgress: bool=True, breakatIndex: int=None, **kwargs):
        """为所有子图应用画图的回调函数，回调函数的参数为 SubplotParams 类型，包含子图信息。
        同时传入的 kwargs 会传递给 add_subplot() 方法。

        Args:
            callback (Callable[[SubplotParams], None]): 回调函数，参数为 SubplotParams 类型。
            **kwargs: 传递给 add_subplot() 方法的参数。
        """
        with Progress(disable=not showProgress) as progress:
            task = progress.add_task("Drawing subplots", 
                                     total=self.nrows*self.ncols \
                                        if breakatIndex is None else breakatIndex+1)

            try:
                for row in range(self.nrows):
                    for col in range(self.ncols):
                        i = row * self.ncols + col
                        colI = row + col * self.nrows
                        ax = self.add_subplot(self.gridSpec[row, col], **kwargs)
                        callback(SubplotParams(row, col, i, colI, ax, self, 
                                            self.nrows, self.ncols,))
                        progress.update(task, advance=1)
                        if breakatIndex is not None and i == breakatIndex:
                            raise ExitLoop()
            except ExitLoop:
                pass

    
    def apply_all_geo_subplots(self, callback: Callable[[SubplotParams], None], 
                               projection: Projection, showProgress: bool=True,
                               breakatIndex: int=None, **kwargs):
        """为所有子图应用画图的回调函数，回调函数的参数为 SubplotParams 类型，包含子图信息，
        mapTools 为 MapTools 类型，为地图工具。
        需要设置 projection 参数，用于设置地图投影。

        Args:
            callback (Callable[[SubplotParams], None]): 回调函数，参数为 SubplotParams 类型。
            **kwargs: 传递给 add_subplot() 方法的参数。
        """
        with Progress(disable=not showProgress) as progress:
            task = progress.add_task("Drawing subplots", 
                                     total=self.nrows*self.ncols \
                                        if breakatIndex is None else breakatIndex+1)

            try:
                for row in range(self.nrows):
                    for col in range(self.ncols):
                        i = row * self.ncols + col
                        colI = row + col * self.nrows
                        ax = self.add_subplot(self.gridSpec[row, col], 
                                            projection=projection, **kwargs)
                        callback(SubplotParams(row, col, i, colI, ax, self, 
                                               self.nrows, self.ncols,
                                               mapTools=MapTools(ax, MapConfig()),
                                               mplTools=MplTools(ax)))
                        progress.update(task, advance=1)
                        if breakatIndex is not None and i == breakatIndex:
                            raise ExitLoop()
            except ExitLoop:
                pass


    
    def adjust_subplots_ratio_by_map_extent(self, which: int=0):
        """
        根据地图的范围比例来重新设置整个图的高度，使得图与地图的比例一致。
        请在定义地图范围后调用。

        Args:
            which (int, optional): 用于参考的子图索引. Defaults to 0.
        """
        extent = self.axes[which].get_extent()
        ratio = (extent[1] - extent[0]) / (extent[3] - extent[2])

        if self.useA4Size:
            figWidth = 8.27
            # 令 subplotWidthInch 变为被动
            subplotWidthInch = (figWidth - self._pt2inch(self.leftPt + self.rightPt) - \
                self._pt2inch(self.wspacePt) * (self.ncols-1)) / self.ncols
        else:
            subplotWidthInch = self.subplotBaseWidthInch

        figHeight = subplotWidthInch / ratio * self.nrows + \
                    self._pt2inch(self.topPt + self.bottomPt) + \
                    self._pt2inch(self.hspacePt) * (self.nrows-1)

        self.set_size_inches(self.get_figwidth(), figHeight)

        self.gridSpec.update(
            hspace=self._pt2inch(self.hspacePt) / (self.subplotBaseWidthInch / ratio),
            top=1-self._pt2inch(self.topPt)/figHeight,
            bottom=self._pt2inch(self.bottomPt)/figHeight
        )

    
    def fig_title(self, title: str, paddingPt: float=5, va='bottom',
                   **kwargs):
        """设置整个图的标题，paddingPt 为标题与最上的子图的间距

        Args:
            title (str): 标题
        """

        subplotTop = self.axes[0].get_position().get_points()[1][1]
        padding = self._pt2inch(paddingPt) / self.get_figheight()
        y = subplotTop + padding

        self.suptitle(title, y=y, va=va, **kwargs)

    
    def fig_hline(self, rows: int | list[int], offsetPt=0, 
                   leftPaddingPt: float=None, rightPaddingPt: float=None, **kwargs):
        """为整个图添加横向分割线，需要首先有网格布局，即调用了 ratio_gridspec() 方法，
        会在第 rows 行的上方添加一条横线。rows 参数接受一个整数或者一个整数列表。

        ！需要在 adjust_subplots_ratio_by_map_extent() 前调用

        Args:
            rows (int | list[int]): 行号
            offsetPt (int, optional): 横线距离下方的偏移，单位为点. Defaults to 0.
            leftPaddingPt (float, optional): 左边距，默认与画幅的左边距相同. Defaults to None.
            rightPaddingPt (float, optional): 右边距，默认与画幅的右边距相同. Defaults to None.
        """
        from matplotlib.lines import Line2D
        if isinstance(rows, int):
            rows = [rows]

        figHeight = self.get_figheight()
        figWidth = self.get_figwidth()
        offset = self._pt2inch(offsetPt) / figHeight

        if leftPaddingPt is None:
            leftPaddingPt = self.leftPt 

        if rightPaddingPt is None:
            rightPaddingPt = self.rightPt

        leftPadding = self._pt2inch(leftPaddingPt) / figWidth
        rightPadding = self._pt2inch(rightPaddingPt) / figWidth

        for row in rows:
            y = self.gridSpec[row, 0].get_position(self).get_points()[1][1]
            self.add_artist(Line2D([leftPadding, 1-rightPadding], 
                                   [y+offset, y+offset], 
                                   transform=self.transFigure, **kwargs))



from xarray import DataArray

class DataScaler:
    def __init__(self, data: ndarray | DataArray, 
                 mode: Literal['maxmin', 'percent']='maxmin', percent: int=95,
                 levelNumber: int=15, minLevelNumber: int=9, powerOfTen: int|None=None,
                 levels: ndarray|None=None,
                 ignoreZero: bool=False, zeroSymmetric: bool=False):
        """数据量级缩放器，将数据尽量缩放到个位数或者十位数，方便绘图

        Args:
            data (ndarray | DataArray): 数据
            mode ("maxmin" | "percent", optional): 计算缩放使用最值还是百分位数. Defaults to 'maxmin'.
            percent (int, optional): 百分位. Defaults to 95.
            levelNumber (int, optional): 等值线数量. Defaults to 10.
            minLevelNumber (int, optional): 等值线最少数量. Defaults to 5.
            powerOfTen (int, optional): 指定数量级. Defaults to 0.
            ignoreZero (bool, optional): 是否去掉零等值线. Defaults to False.

        Raises:
            ValueError: mode必须是'maxmin'或'percent'
        """
        self.data = data
        self.scale_factor: int = 1  # 缩放倍数
        self.powerOfTen: int | None = powerOfTen    # 数量级

        self.vMax = None
        self.vMin = None

        self.levels: ndarray | None = None
        self.ignoreZero = ignoreZero
        self.zeroSymmetric = zeroSymmetric

        if mode == 'maxmin':
            self.vMax = np.nanmax(self.data)
            self.vMin = np.nanmin(self.data)

        elif mode == 'percent':
            if percent >= 100 or percent <= 0:
                raise ValueError("percent must be in (0, 100)")

            if percent < 50:
                percent = 100 - percent            

            self.vMax = np.nanpercentile(self.data, percent)
            self.vMin = np.nanpercentile(self.data, 100 - percent)

        else:
            raise ValueError("mode must be 'maxmin' or 'percent'")

        if levels is not None:
            self.levels = levels
            if powerOfTen is not None:
                self.powerOfTen = -powerOfTen
                self.scale_factor = 10 ** -powerOfTen
            else:
                self.powerOfTen = self._get_power_of_ten(self.vMax, self.vMin)
                self.scale_factor = 10 ** -self.powerOfTen
        else:
            self.levels = self._getLevels(self.vMax, self.vMin, levelNumber, minLevelNumber)

        if ignoreZero:
            self.levels = self.levels[self.levels != 0]
    

    def _get_power_of_ten(self, vMax: float, vMin: float) -> int:
        vMaxAbs = np.abs(vMax)
        vMinAbs = np.abs(vMin)
        maxN = np.log10(vMaxAbs)
        minN = np.log10(vMinAbs)

        # 取绝对值最大的数的数量级
        scaleN = np.floor(np.max([maxN, minN]))
        factor = 10 ** -scaleN

        # 如果最大值和最小值的值小于3，则将数量级减小1，数据缩放到-30到30之间
        if np.max([vMaxAbs * factor, vMinAbs * factor]) < 3:
            scaleN -= 1
            factor = 10 ** -scaleN
        
        return scaleN


    def _getLevels(self, vMax: float , vMin: float, levelNumer: int=10, 
                   minLevelNumber: int=5) -> ndarray:
        """计算等值线的数量级

        Args:
            vMax (number): 最大值
            vMin (number): 最小值
            levelNumer (int, optional): 等值线数量. Defaults to 10.
            minLevelNumber (int, optional): 等值线最少数量. Defaults to 5.

        Returns:
            ndarray: 等值线数组
        """
        from matplotlib.ticker import MaxNLocator

        vMaxAbs = np.abs(vMax)
        vMinAbs = np.abs(vMin)
        maxN = np.log10(vMaxAbs)
        minN = np.log10(vMinAbs)

        # 如果指定了数量级，则直接返回
        if self.powerOfTen is not None:
            scaleN = self.powerOfTen
            factor = 10 ** -scaleN
            self.scale_factor = factor
            levels = MaxNLocator(nbins=levelNumer, integer=True, 
                                 steps=[1, 2, 5], symmetric=self.zeroSymmetric)\
                .tick_values(vMin * factor, vMax * factor)
            return levels

        # 取绝对值最大的数的数量级
        scaleN = np.floor(np.max([maxN, minN]))
        factor = 10 ** -scaleN

        # 如果最大值和最小值的值小于3，则将数量级减小1，数据缩放到-30到30之间
        if np.max([vMaxAbs * factor, vMinAbs * factor]) < 3:
            scaleN -= 1
            factor = 10 ** -scaleN
        
        levels = MaxNLocator(nbins=levelNumer, integer=True, 
                             steps=[1, 2, 5], symmetric=self.zeroSymmetric)\
            .tick_values(vMin * factor, vMax * factor)
        
        if len(levels) < minLevelNumber or levels[-1] - levels[0] < 3:
            scaleN -= 1
            factor = 10 ** -scaleN
            levels = MaxNLocator(nbins=10, integer=True, 
                                 steps=[1, 2, 5], symmetric=self.zeroSymmetric)\
                .tick_values(vMin * factor, vMax * factor)
        
        self.powerOfTen = scaleN
        self.scale_factor = factor
        
        return levels


    def applyScale(self, data: DataArray | ndarray) -> DataArray | ndarray:
        """应用缩放到数据

        Args:
            data (DataArray | ndarray): 数据

        Returns:
            DataArray | ndarray: 缩放后的数据
        """
        return data * self.scale_factor

    def latex_unit(self, unitLatex: str):
        if self.powerOfTen == 0:
            return f"${unitLatex}$"

        if self.powerOfTen == 1:
            return f"$10\cdot {unitLatex}$"

        return f"$10^{{{self.powerOfTen:.0f}}}\cdot {unitLatex}$"